import pymysql
from database_manager import get_db_connection

def update_industry_accurate():
    """
    更精确地更新股票行业信息
    使用详细的股票代码前缀映射关系
    """
    conn = get_db_connection()
    if not conn:
        print("无法连接数据库")
        return
    
    try:
        with conn.cursor() as cursor:
            # 创建更详细的行业映射字典
            # 基于股票代码前缀和交易所分类进行更精确的行业映射
            industry_mapping = {
                # 上证指数(600开头)
                '600000': '银行-浦发银行',
                '600016': '银行-民生银行',
                '600019': '钢铁-宝钢股份',
                '600028': '石油-中国石化',
                '600036': '银行-招商银行',
                '600050': '通信-中国联通',
                '600104': '汽车-上汽集团',
                '600111': '有色金属-包钢稀土',
                '600519': '白酒-贵州茅台',
                '600887': '食品饮料-伊利股份',
                '601088': '煤炭-中国神华',
                '601166': '银行-兴业银行',
                '601288': '银行-农业银行',
                '601318': '保险-中国平安',
                '601328': '银行-交通银行',
                '601398': '银行-工商银行',
                '601628': '保险-中国人寿',
                '601668': '建筑-中国建筑',
                '601688': '证券-华泰证券',
                '601857': '石油-中国石油',
                '601939': '银行-建设银行',
                '601988': '银行-中国银行',
                '601998': '银行-中信银行',
                
                # 深证成指(000开头)
                '000002': '房地产-万科A',
                '000001': '银行-平安银行',
                '000063': '电子-中兴通讯',
                '000333': '家电-美的集团',
                '000568': '白酒-泸州老窖',
                '000651': '家电-格力电器',
                '000725': '电子-京东方A',
                '000858': '白酒-五粮液',
                '000895': '医药-双汇发展',
                '000938': '煤炭-紫光股份',
                
                # 中小板(002开头)
                '002024': '家电-苏宁易购',
                '002230': '电子-科大讯飞',
                '002352': '汽车-顺丰控股',
                '002415': '医药-海康威视',
                '002594': '新能源汽车-比亚迪',
                
                # 创业板(300开头)
                '300059': '医药-东方财富',
                '300750': '医药-宁德时代',
                '300015': '医药-爱尔眼科',
                
                # 科创板(688开头)
                '688981': '医药-中芯国际',
                '688005': '科技-容百科技',
                '688036': '科技-传音控股',
            }
            
            # 首先按精确匹配更新已知股票
            updated_count = 0
            for code, industry in industry_mapping.items():
                cursor.execute(
                    "UPDATE stock_info SET industry = %s WHERE symbol_standardized = %s",
                    (industry, code)
                )
                updated_count += cursor.rowcount
            
            print(f"已精确更新 {updated_count} 条股票的行业信息")
            
            # 然后按前缀进行更细致的分类
            prefix_mapping = {
                # 银行类
                '60000': '银行',
                '60001': '银行',
                '60003': '银行',
                '60116': '银行',
                '60128': '银行',
                '60132': '银行',
                '60139': '银行',
                '60193': '银行',
                '60198': '银行',
                '60199': '银行',
                '000001': '银行',
                
                # 保险类
                '60131': '保险',
                '60162': '保险',
                
                # 证券类
                '60168': '证券',
                '60178': '证券',
                '60188': '证券',
                
                # 医药类
                '00051': '医药',
                '00053': '医药',
                '00200': '医药',
                '00203': '医药',
                '00222': '医药',
                '30000': '医药',
                '30001': '医药',
                '30003': '医药',
                '30005': '医药',
                
                # 科技电子类
                '00006': '科技',
                '00072': '科技',
                '00073': '科技',
                '00079': '科技',
                '00223': '科技',
                '00237': '科技',
                '30002': '科技',
                '30013': '科技',
                '60053': '科技',
                '60074': '科技',
                '60080': '科技',
                
                # 白酒食品饮料
                '00056': '白酒',
                '00059': '白酒',
                '00085': '白酒',
                '60051': '白酒',
                '60077': '白酒',
                '60336': '白酒',
                '60358': '白酒',
                '60088': '食品饮料',
                '00089': '食品饮料',
                '00250': '食品饮料',
                
                # 新能源汽车
                '00259': '新能源汽车',
                '00270': '新能源汽车',
                '60010': '新能源汽车',
                '60041': '新能源汽车',
                '60073': '新能源汽车',
                
                # 房地产
                '00000': '房地产',
                '00004': '房地产',
                '00006': '房地产',
                '00040': '房地产',
                '00060': '房地产',
                '00067': '房地产',
                '00071': '房地产',
                '00092': '房地产',
                '00214': '房地产',
                '60004': '房地产',
                '60006': '房地产',
                '60009': '房地产',
                '60017': '房地产',
                '60020': '房地产',
                
                # 家电
                '00033': '家电',
                '00065': '家电',
                '00202': '家电',
                '00224': '家电',
                '60069': '家电',
                '60083': '家电',
                
                # 煤炭石油
                '60002': '石油',
                '60018': '煤炭',
                '60034': '煤炭',
                '60039': '煤炭',
                '60050': '煤炭',
                '60108': '煤炭',
                '60185': '石油',
                
                # 钢铁
                '60001': '钢铁',
                '60002': '钢铁',
                '60011': '钢铁',
                '60028': '钢铁',
                '60030': '钢铁',
                '60035': '钢铁',
                '60056': '钢铁',
                
                # 建筑
                '60017': '建筑',
                '60026': '建筑',
                '60085': '建筑',
                '60166': '建筑',
                '60176': '建筑',
                
                # 通信
                '60005': '通信',
                '00006': '通信',
                '00083': '通信',
                '00093': '通信',
                
                # 汽车
                '00062': '汽车',
                '00080': '汽车',
                '00092': '汽车',
                '00204': '汽车',
                '00235': '汽车',
                '60010': '汽车',
                '60033': '汽车',
                '60046': '汽车',
            }
            
            # 更新剩余的股票，按前缀分类
            for prefix, industry in prefix_mapping.items():
                # 只更新industry为空或为'其他'的记录
                cursor.execute(
                    "UPDATE stock_info SET industry = %s WHERE symbol_standardized LIKE %s AND (industry IS NULL OR industry = '其他')",
                    (industry, prefix + '%')
                )
            
            # 最后，对于仍然没有行业信息的股票，使用更通用的分类
            sql = """
                UPDATE stock_info SET industry = CASE 
                    WHEN symbol_standardized LIKE '600%' THEN '主板-工业'
                    WHEN symbol_standardized LIKE '601%' THEN '主板-金融'
                    WHEN symbol_standardized LIKE '603%' THEN '主板-消费'
                    WHEN symbol_standardized LIKE '000%' THEN '深市主板'
                    WHEN symbol_standardized LIKE '001%' THEN '深市主板'
                    WHEN symbol_standardized LIKE '002%' THEN '中小板'
                    WHEN symbol_standardized LIKE '300%' THEN '创业板'
                    WHEN symbol_standardized LIKE '688%' THEN '科创板'
                    ELSE '其他行业'
                END WHERE industry IS NULL OR industry = '其他' OR industry = '其他'
            """
            cursor.execute(sql)
            
            # 获取最终更新后的行业分布
            cursor.execute("SELECT industry, COUNT(*) as count FROM stock_info GROUP BY industry ORDER BY count DESC")
            industry_distribution = cursor.fetchall()
            
            print("\n行业分布情况：")
            for item in industry_distribution[:20]:  # 只显示前20个行业
                print(f"{item['industry']}: {item['count']} 条")
            
            # 获取更新后的总数
            cursor.execute("SELECT COUNT(*) FROM stock_info WHERE industry IS NOT NULL AND industry != '其他' AND industry != '其他'")
            valid_count = cursor.fetchone()['COUNT(*)']
            print(f"\n有效行业信息的股票数量: {valid_count}")
            
            # 提交事务
            conn.commit()
            print("\n行业信息更新完成")
            
    except Exception as e:
        print(f"更新行业信息时出错: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    update_industry_accurate()
